import math
import torch
from torch import nn
import torch.nn.functional as F
import numpy as np
import train_utils.distributed_utils as utils
from .dice_coefficient_loss import dice_loss, build_target


def criterion(inputs, target, loss_weight=None, num_classes: int = 2, dice: bool = True, ignore_index: int = -100):        #计算loss和dice
    losses = {}
    ce_loss_weight = 1.0  # 交叉熵损失权重
    dice_loss_weight = 2.0  # Dice损失权重

    for name, x in inputs.items():                                                                       #inputs.items() 返回一个可迭代的视图对象，该视图对象包含了字典 inputs 中所有的键值对。
        # ignore_index：忽略target中值为255的像素，255的像素是目标边缘或者padding填充
        ce_loss = nn.functional.cross_entropy(x, target, ignore_index=ignore_index, weight=loss_weight)    #计算了给定预测张量 x 和目标标签张量 target 之间的交叉熵损失。
        if dice is True:
            dice_target = build_target(target, num_classes, ignore_index)                                                  #对不同的类别分别构造GT
            d_loss = dice_loss(x, dice_target, multiclass=True, ignore_index=ignore_index)
            loss = ce_loss * ce_loss_weight + d_loss * dice_loss_weight

        losses[name] = loss

    if len(losses) == 1:
        return losses['out']

    return losses['out'] + 0.5 * losses['aux']


def evaluate(model, data_loader, device, num_classes):
    model.eval()
    confmat = utils.ConfusionMatrix(num_classes)
    dice = utils.DiceCoefficient(num_classes=num_classes, ignore_index=255)
    metric_logger = utils.MetricLogger(delimiter="  ")                                         #用于创建一个度量日志记录器（Metric Logger）的对象(用于跟踪和记录模型训练过程中的各种度量指标，如损失值、准确率、学习率等)，并设置了分隔符（delimiter）为两个空格。
    header = 'Test:'
    with torch.no_grad():                                                                                         #torch.no_grad() 是 PyTorch 中的一个上下文管理器（Context Manager），用于指定在其内部的代码块中不需要计算梯度。
        for image, target in metric_logger.log_every(data_loader, 100, header):   #通常是自定义的工具类或模块，用于简化训练过程中的日志记录和指标监控。
            image, target = image.to(device), target.to(device)
            output = model(image)
            output = output['out']                                                                                #用于从模型输出的字典中提取特定键 'out' 对应的值，并将其赋给变量 output

            confmat.update(target.flatten(), output.argmax(1).flatten())
            dice.update(output, target)

        confmat.reduce_from_all_processes()
        dice.reduce_from_all_processes()

    return confmat, dice.value.item()


def entropy(tensor):
    """
    计算标准化到 [-1,1] 范围的灰度图像张量的熵。
    参数:
    tensor_img: 输入的灰度图像张量，类型为 torch.Tensor。
    返回值:
    entropy: 计算得到的熵值。
    """
    # 将张量映射回 [0,1] 范围
    original_img = (tensor + 1) / 2
    # 将图像像素值缩放到 [0,255] 范围，并转换为整数类型
    scaled_img = (original_img * 255).to(torch.uint8)
    # 计算灰度图像的直方图
    histogram = torch.histc(scaled_img.float(), bins=256, min=0, max=255)
    # 计算概率分布
    probabilities = histogram / torch.sum(histogram + 1e-10)
    # 计算熵
    entropy = -torch.sum(probabilities * torch.log2(probabilities + 1e-10))

    return entropy.item()  # 返回熵值


def addimages(model, data_loader, device, num_classes):
    model.eval()
    count=1
    entropy_dict = {}
    with torch.no_grad():                                                                                         #torch.no_grad() 是 PyTorch 中的一个上下文管理器（Context Manager），用于指定在其内部的代码块中不需要计算梯度。
        for image, target in data_loader:                                                                  #torch.utils.data.DataLoader() 默认按照文件系统返回的顺序加载文件，而文件系统通常按照字符串字典序返回文件名列表。因此，对于文件名为 "1, 2, 3, ..., 10, 11, 12" 的情况，DataLoader() 将按照这些文件名的字典序来加载数据，即按照 "1, 10, 11, 12, 2, 3, ..." 的顺序。
            image, target = image.to(device), target.to(device)
            output = model(image)
            output = output['out']
            entropy_values = entropy(output)

            entropy_dict[count] = entropy_values
            count += 1
    max_entropy = max(entropy_dict, key=entropy_dict.get)
    return max_entropy


def train_one_epoch(model, optimizer, data_loader, device, epoch, num_classes,
                    lr_scheduler, print_freq=10, scaler=None,pause_event=None):
    model.train()
    metric_logger = utils.MetricLogger(delimiter="  ")
    metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
    header = 'Epoch: [{}]'.format(epoch)

    if num_classes == 2:
        # 设置cross_entropy中背景和前景的loss权重(根据自己的数据集进行设置：前景占的比例越小则前景的权重越大)
        loss_weight = torch.as_tensor([1.0, 3.0], device=device)
    else:
        loss_weight = None

    for image, target in metric_logger.log_every(data_loader, print_freq, header):
        pause_event.wait()
        image, target = image.to(device), target.to(device)
        with torch.cuda.amp.autocast(enabled=scaler is not None):
            output = model(image)
            loss = criterion(output, target, loss_weight, num_classes=num_classes, ignore_index=255)

        optimizer.zero_grad()                                                                            #将当前梯度清空，避免和上一次累加
        if scaler is not None:                                                                              #使不使用梯度缩放器
            scaler.scale(loss).backward()                                                             #算梯度
            scaler.step(optimizer)                                                                        #沿着梯度方向更新
            scaler.update()                                                                                    #更新放大因子
        else:
            loss.backward()
            optimizer.step()

        lr_scheduler.step()                                                                                 #更新学习率调度器状态

        lr = optimizer.param_groups[0]["lr"]
        metric_logger.update(loss=loss.item(), lr=lr)

    return metric_logger.meters["loss"].global_avg, lr


def create_lr_scheduler(optimizer,
                        num_step: int,
                        epochs: int,
                        warmup=True,
                        warmup_epochs=1,
                        warmup_factor=1e-3):
    assert num_step > 0 and epochs > 0
    if warmup is False:
        warmup_epochs = 0

    #def cosine_annealing(x, T_max, eta_min=0):
        #return eta_min + (0.5 * (1 + math.cos(x / T_max * math.pi)))

    def f(x):
        """
        根据step数返回一个学习率倍率因子，
        注意在训练开始之前，pytorch会提前调用一次lr_scheduler.step()方法
        """
        if warmup is True and x <= (warmup_epochs * num_step):
            alpha = float(x) / (warmup_epochs * num_step)
            # warmup过程中lr倍率因子从warmup_factor -> 1
            return warmup_factor * (1 - alpha) + alpha
        else:
            #余弦退火
            #T_max = (epochs - warmup_epochs) * num_step
            #return cosine_annealing(x - warmup_epochs * num_step, T_max)
            # warmup后lr倍率因子从1 -> 0
            # 参考deeplab_v2: Learning rate policy
            return (1 - (x - warmup_epochs * num_step) / ((epochs - warmup_epochs) * num_step)) ** 0.9

    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=f)
